
#include <optix_cuda.h>
#include <optix_math.h>

#include "helpers.h"
#include "random.h"
#include "OptixData.h"

#define POINT_LIGHT 0
#define AREA_LIGHT 3

struct PerRayData_radiance
{
	float3 position;
	float3 radiance;
	uint2 seed;
	uint index;
	int	bounce;
};

rtDeclareVariable(float3,				bad_color, , );
rtDeclareVariable(float3,				bg_color, , );
rtDeclareVariable(float,				scene_epsilon, , );

//this is the object all traces will start with
rtDeclareVariable(rtObject,				top_object, , );

//this is the maximum number of bounces a light photon will take
rtDeclareVariable(unsigned int,			MaxBounces, , );
rtDeclareVariable(OptixLight,			Light, , );
rtDeclareVariable(float,				IndirectIntensity, , );

rtDeclareVariable(uint,					radiance_ray_type, , );

//input buffers
rtBuffer< uint2, 2 >					seed_buffer;

//output buffers
rtBuffer< VirtualPointLight, 1 >		output_vpls;

//optix tracked vars
rtDeclareVariable(optix::Ray, ray,		rtCurrentRay, );
rtDeclareVariable(uint2, launch_index,	rtLaunchIndex, );
rtDeclareVariable(uint2, launch_dim,	rtLaunchDim, );
rtDeclareVariable(float, t_hit,			rtIntersectionDistance, );

//ray data
rtDeclareVariable(PerRayData_radiance, prd_radiance, rtPayload, );

__device__ __inline__ float3 generateHemisphereLightPhoton( const float2& sample, const float3& direction )
{
	// Create an ortho-normal basis about the direction
	float3 U, V, W;
	createONB( direction, U, V, W );

	//sample a unit hemisphere about the basis
	return sampleUnitHemisphere( sample, U, V, W );
}

RT_PROGRAM void instant_radiosity()
{
	//get our random seed
	uint2 seed = seed_buffer[ launch_index ];

	//create a random photon direction
	float2 raySeed = make_float2( ( (float)launch_index.x + rnd( seed.x ) ) / (float)launch_dim.x,
								  ( (float)launch_index.y + rnd( seed.y ) ) / (float)launch_dim.y );

	//Note: to be phsyically accurate we should be generating photons in the unit sphere around the light for the
	//sponza demo since it uses a point light not an area light.
	float3 origin = Light.position;
	float3 direction = generateHemisphereLightPhoton( raySeed, Light.direction );

	//create our ray
	optix::Ray ray(origin, direction, radiance_ray_type, scene_epsilon );

	//create our ray data packet and launch a ray
	PerRayData_radiance prd;
	prd.radiance = Light.color * Light.intensity * IndirectIntensity;
	prd.bounce = 0;
	prd.seed = seed;
	prd.index = ( launch_index.y * launch_dim.x + launch_index.x ) * MaxBounces;
	rtTrace( top_object, ray, prd );
}

//surface attributes
rtDeclareVariable(float3, texcoord,			attribute texcoord, ); 
rtDeclareVariable(float3, shading_normal,	attribute shading_normal, );
rtDeclareVariable(float3, geometric_normal, attribute geometric_normal, ); 

rtDeclareVariable(float3, diffuse_color, , );
//rtTextureSampler<uchar4, 2, cudaReadModeNormalizedFloat> diffuseTex;
RT_PROGRAM void closest_hit_radiosity()
{
	//convert the geometry's normal to world space
	//RT_OBJECT_TO_WORLD is an Optix provided transformation
	float3 world_shading_normal   = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, shading_normal ) );
	float3 world_geometric_normal = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, geometric_normal ) );
	float3 ffnormal     = world_shading_normal;//faceforward( world_shading_normal, -ray.direction, world_geometric_normal );

	//calculate the hitpoint of the ray
	float3 hit_point = ray.origin + t_hit * ray.direction;

	//sample the texture for the geometry
	float3 Kd = diffuse_color;//norm_rgb( tex2D( diffuseTex, texcoord.x, texcoord.y ) );
	//Kd = pow3f( Kd, 2.2f ); //convert to linear space
	//Kd *= make_float3( diffuseColor ); //multiply the diffuse material color
	prd_radiance.radiance = Kd * prd_radiance.radiance; //calculate the ray's new radiance value

	// We hit a diffuse surface; record hit if it has bounced at least once
	if( prd_radiance.bounce >= 0 ) 
	{
		//offset the light a bit from the hit point
		float3 lightPos = ray.origin + ( t_hit - 0.1f ) * ray.direction;

		VirtualPointLight vpl;
		vpl.position = make_float4( lightPos, t_hit * 0.5f );

		//the light's intensity is divided equally among the photons. Each photon starts out with an intensity
		//equal to the light. So here we must divide by the number of photons cast from the light.
		vpl.radiance = make_float4( prd_radiance.radiance * 1.0f / ( launch_dim.x * launch_dim.y ), 1.0f );

		 output_vpls[ prd_radiance.index + prd_radiance.bounce ] = vpl;
	}

	//if we're less than the max number of bounces shoot another ray
	//we could also implement Russion Roulette here so that we would have a less biased solution
	prd_radiance.bounce++;
	if ( prd_radiance.bounce >= MaxBounces )
		return;

	//here we "rotate" the seeds in order to have a little more variance
	prd_radiance.seed.x = prd_radiance.seed.x ^ prd_radiance.bounce;
	prd_radiance.seed.y = prd_radiance.seed.y ^ prd_radiance.bounce;
	float2 seed_direction = make_float2( ( (float)launch_index.x + rnd( prd_radiance.seed.x ) ) / (float)launch_dim.x,
										 ( (float)launch_index.y + rnd( prd_radiance.seed.y ) ) / (float)launch_dim.y );

	//generate a new ray in the hemisphere oriented to the surface
	float3 new_ray_dir = generateHemisphereLightPhoton( seed_direction, ffnormal );

	//cast a new ray into the scene
	optix::Ray new_ray( hit_point, new_ray_dir, radiance_ray_type, scene_epsilon );
	rtTrace(top_object, new_ray, prd_radiance);
}

RT_PROGRAM void exception()
{
	rtPrintExceptionDetails();
}

RT_PROGRAM void miss()
{
	prd_radiance.radiance = make_float3( .39f, .58f, .93f ) * prd_radiance.radiance; //calculate the ray's new radiance value

	//offset the light a bit from the hit point
	float3 lightPos = ray.origin + ( 15.0f ) * ray.direction;

	VirtualPointLight vpl;
	vpl.position = make_float4( lightPos, 5.0f * 0.5f );

	//the light's intensity is divided equally among the photons. Each photon starts out with an intensity
	//equal to the light. So here we must divide by the number of photons cast from the light.
	vpl.radiance = make_float4( prd_radiance.radiance * 1.0f / ( launch_dim.x * launch_dim.y ), 1.0f );

	 output_vpls[ prd_radiance.index + prd_radiance.bounce ] = vpl;
	//1.0f, .39f, .58f, .93f
}